# 一个双均线策略回测
import akshare as ak
import numpy as np
import datetime
import matplotlib.pyplot as plt
import pandas as pd

from datacache import get_stock_data


def calculate_moving_averages(df, short_window, long_window):
    """
    计算短期和长期移动平均线
    :param df: 股票数据
    :param short_window: 短期窗口
    :param long_window: 长期窗口
    :return: 包含移动平均线的DataFrame
    """
    df['ma_short'] = df['close'].rolling(window=short_window, min_periods=1).mean()
    df['ma_long'] = df['close'].rolling(window=long_window, min_periods=1).mean()
    return df

def generate_signals(df):
    """
    生成交易信号
    :param df: 包含移动平均线的DataFrame
    :return: 包含信号的DataFrame
    """
    df['signal'] = 0
    # 当短期均线上穿长期均线时，生成买入信号
    df['signal'][df['ma_short'] > df['ma_long']] = 1
    # 当短期均线下穿长期均线时，生成卖出信号
    df['signal'][df['ma_short'] < df['ma_long']] = -1
    # 计算信号的变化点
    df['positions'] = df['signal'].diff()
    return df

def backtest_strategy(df):
    """
    回测策略
    :param df: 包含信号的DataFrame
    :return: 回测结果的DataFrame
    """
    # 初始化资金
    initial_capital = 100000  # 初始资金10万元
    positions = pd.DataFrame(index=df.index).fillna(0.0)
    # 假设每次交易买入100股
    positions['stock'] = df['signal'] * 100
    # 计算每日持仓的价值
    portfolio = positions.multiply(df['close'], axis=0)
    # 计算现金
    pos_diff = positions['stock'].diff()
    # 买入时减少现金，卖出时增加现金
    portfolio['cash'] = initial_capital - (pos_diff.multiply(df['close'], axis=0)).cumsum()
    # 总资产
    portfolio['total'] = portfolio['stock'].multiply(df['close'], axis=0) + portfolio['cash']
    # 计算收益
    portfolio['returns'] = portfolio['total'].pct_change()
    return portfolio

def plot_strategy(df, portfolio):
    """
    绘制策略图表
    :param df: 包含信号的DataFrame
    :param portfolio: 回测结果的DataFrame
    """
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 10), sharex=True)

    # 绘制价格和移动平均线
    ax1.plot(df.index, df['close'], label='收盘价', color='blue')
    ax1.plot(df.index, df['ma_short'], label='短期均线', color='red')
    ax1.plot(df.index, df['ma_long'], label='长期均线', color='green')
    # 绘制买入信号
    ax1.plot(df.index[df.positions == 1], df['ma_short'][df.positions == 1], '^', markersize=10, color='m', label='买入信号')
    # 绘制卖出信号
    ax1.plot(df.index[df.positions == -1], df['ma_short'][df.positions == -1], 'v', markersize=10, color='k', label='卖出信号')
    ax1.set_title('双均线策略回测')
    ax1.set_ylabel('价格')
    ax1.legend()

    # 绘制总资产
    ax2.plot(portfolio.index, portfolio['total'], label='总资产', color='b')
    ax2.set_ylabel('总资产 (元)')
    ax2.legend()

    plt.xlabel('日期')
    plt.show()

def main():
    # 参数设置
    stock_code = 'sz000001'  # 平安银行
    start_date = '20240101'
    end_date = '20250101'
    short_window = 50
    long_window = 200

    # 获取数据
    df = get_stock_data(stock_code, start_date, end_date)
    print(f"获取到的数据日期范围: {df.index.min()} 到 {df.index.max()}")

    # 计算移动平均线
    df = calculate_moving_averages(df, short_window, long_window)

    # 生成交易信号
    df = generate_signals(df)

    # 回测策略
    portfolio = backtest_strategy(df)

    # 输出回测结果
    total_return = (portfolio['total'][-1] / portfolio['total'][0]) - 1
    print(f"初始资金: {portfolio['total'][0]}元")
    print(f"最终资金: {portfolio['total'][-1]}元")
    print(f"总收益率: {total_return * 100:.2f}%")

    # 绘制图表
    plot_strategy(df, portfolio)

if __name__ == "__main__":
    main()